//Examples/reed_sol_crs_03.cpp

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "jerasure.h"
#include "reed_sol.h"

#define talloc(type, num) (type *) malloc(sizeof(type)*(num))

void usage(char *s)
{
	fprintf(stderr, "usage: reed_sol_crs_03 k m w - Scheduled Cauchy Reed-Solomon coding example in GF(2^w).\n");
	fprintf(stderr, "       \n");
	fprintf(stderr, "       k+m must be <= 2^w.  It sets up a Cauchy distribution matrix and encodes\n");
	fprintf(stderr, "       k sets of w*8 bytes. It encodes and decodes using both smart and dumb bit-matrix scheduling.\n");
	fprintf(stderr, "       \n");
	fprintf(stderr, "This demonstrates: CRS_Generator()\n");
	fprintf(stderr, "                   CRS_Convert_To_Bitmatrix()\n");
	fprintf(stderr, "                   JER_Matrix::Print()\n");
	fprintf(stderr, "                   JER_Gen_T::Create_Encode_Schedule()\n");
	fprintf(stderr, "                   JER_Slices::Encode()\n");
	fprintf(stderr, "                   JER_Slices::Add_Drive_Failure()\n");
	fprintf(stderr, "                   JER_Slices::Decode_Schedule_Lazy()\n");
	if (s != NULL) fprintf(stderr, "%s\n", s);
	exit(1);
}

static void print_data_and_coding(JER_Slices *slices){

	int i, j, x, y;
	int k, m, w, size;
	int n, sp;
	unsigned char uc;
	//unsigned char l[8]; //holds 1 uint64_t
	uint64_t l;

	JER_Gen_T * g;
	vector < string > states;
	int max_data_state_len; //max strlen of a data drive state
	int max_coding_state_len;
	int state_len;

	states.resize(3);
	states[0] = "up";
	states[1] = "down";
	states[2] = "unusable";

	max_data_state_len = max_coding_state_len = 2;
	for(i=0;i<slices->N;i++){
		state_len = strlen(states[slices->States[i]].c_str());
		if(i<slices->K){
			//data
			if(state_len>max_data_state_len){
				max_data_state_len = state_len;
			}
		}else{
			//coding
			if(state_len>max_coding_state_len){
				max_coding_state_len = state_len;
			}
		}
	}

	g = slices->G;
	if(g == NULL){
		return;
	}

	if(g->M == NULL){
		return;
	}

	k = slices->K;
	m = slices->N - slices->K;
	w = g->WPD;
	size = slices->PacketSize * slices->PacketsPerSlice;

	if(k > m) n = k;
	else n = m;

	//14 base spaces, 2 spaces for each byte in a packet, 3+longest state name
	sp = 13 + slices->PacketSize * 2 + max_data_state_len + 3;

	printf("%-*sCoding\n", sp, "Data");
	for (i = 0; i < n; i++) {
		for (j = 0; j< w; j++) {
			//print one of the 'w' packets
			if (i < k) {

				if (j == 0) printf("D%-2d (%-*s): p%-2d:", i,max_data_state_len,states[slices->States[i]].c_str(), j);
				else        printf("%-*sp%-2d:",8+max_data_state_len,"",j);

				for (x = 0; x < slices->PacketSize; x+=8) {
					//get a packet
					memcpy(&l, slices->Ptrs[i]+j*slices->PacketSize+x, sizeof(uint64_t));
					for (y = 0; y < 8; y++){
						//get a byte from the uint64_t
						//memcpy(&uc,&l+y,sizeof(char));
						uc = (l>>((7-y)*8))&255;
						printf("%02x", uc);
					}
				}
				printf("    ");
			}
			else printf("%*s", sp, "");
			if (i < m) {
				if (j == 0) printf("C%-2d (%-*s): p%-2d:", i, max_coding_state_len,states[slices->States[i+k]].c_str(),j);
				else        printf("%-*sp%-2d:",8+max_coding_state_len,"",j);
				for ( x = 0; x < slices->PacketSize; x+=8) {
					memcpy(&l, slices->Ptrs[k+i] + j*slices->PacketSize + x, sizeof(uint64_t));
					for (y = 0; y < 8; y++){
						//get a byte from the uint64_t
						//memcpy(&uc,&l+y,sizeof(char));
						uc = (l>>((7-y)*8))&255;
						printf("%02x", uc);
					}
				}
			}
			printf("\n");
		}
	}
	printf("\n");
}

int main(int argc, char **argv)
{
	uint64_t l;
	int k, w, i, j, m;
	vector <int> erasures;
	vector <int> erased;
	vector <double> stats;
	JER_Slices * slices;
	JER_Gen_T * g;
	int psize;
	int size;
	int x;
	unsigned char *data;

	if (argc != 4) usage(NULL);
	if (sscanf(argv[1], "%d", &k) == 0 || k <= 0) usage((char*)"Bad k");
	if (sscanf(argv[2], "%d", &m) == 0 || m <= 0) usage((char*)"Bad m");
	if (sscanf(argv[3], "%d", &w) == 0 || w <= 0 || w > 32) usage((char *)"Bad m");
	if (w < 30 && (k+m) > (1 << w)) usage((char *)"k + m is too big");
	psize = 8;

	g = CRS_Generator(k,m,w);
	if(g == NULL){
		printf("Error creating generator\n");
		return 0;	
	}

	if(CRS_Convert_To_Bitmatrix(*g) == -1){
		printf("Error converting generator to bitmatrix\n");
		return 0;	
	}

	printf("Last (m * w) rows of the Binary Distribution Matrix:\n");
	printf("\n");

	g->M->Print(w);

	//create slices:
	//w packets per slice
	slices = new JER_Slices(k+m, k,psize, w,g);

	size = w*psize;
	data = new unsigned char [(k+m) * size];
	bzero(data,(k+m)*size);

	//setup data drives
	//each drive has 'w' packets of size psize
	for (i = 0; i < k; i++) {
		slices->Ptrs[i] = data + i*w*psize;
		for (j = 0; j < w; j++) {
			for (x = 0; x < psize; x += 8) {
				l = lrand48();
				l <<= 8*4;
				l += lrand48();
				memcpy(data + i*psize*w + j*psize + x, &l, 8);
			}
		}
	}

	//setup coding drives
	for (i = 0; i < m; i++) {
		slices->Ptrs[k + i] = data + k*psize*w + i*w*psize;
	}

	printf("\nOriginal data:\n\n");
	print_data_and_coding(slices);
	
	//create dumb schedule for encoding
	g->Create_Encode_Schedule(false);

	//encode with the dumb schedule
	slices->Encode();

	printf("Dumb Schedule Encoding Complete: - %llu XOR'd bytes\n\n", slices->XORs);
	print_data_and_coding(slices);
	slices->XORs = 0;

	//create smart schedule for encoding
	g->Create_Encode_Schedule(true);

	//encode with the smart schedule
	slices->Encode();

	printf("Smart Schedule Encoding Complete: - %llu XOR'd bytes\n\n", slices->XORs);
	print_data_and_coding(slices);
	slices->XORs = 0;

	//erase m drives
	erasures.resize(m,0);
	for (i = 0; i < m; ) {
		erasures[i] = lrand48()%(k+m);
		if (slices->States[erasures[i]] == 0) {
			slices->States[erasures[i]] = 1;
			bzero(slices->Ptrs[erasures[i]], size);
			i++;
		}
	}

	printf("Erased %d random devices:\n\n", m);
	print_data_and_coding(slices);

	//decode with a smart schedule
	slices->Decode_Schedule_Lazy(true);	

	printf("Decoded using a smart schedule - %llu XOR'd bytes\n\n", slices->XORs);
	print_data_and_coding(slices);
	slices->XORs = 0;

	//erase the same devices
	for (i = 0; i < m; ) {
		slices->States[erasures[i]] = 1;
		bzero(slices->Ptrs[erasures[i]], size);
		i++;
	}

	printf("Erased the same %d devices:\n\n", m);
	print_data_and_coding(slices);


	//decode with a dumb schedule
	slices->Decode_Schedule_Lazy(false);	
	printf("Decoded using a dumb schedule - %llu XOR'd bytes\n\n", slices->XORs);
	print_data_and_coding(slices);
	slices->XORs = 0;

	delete slices;
	delete g;

	return 0;
}
